# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

# Define input data and filter
input_data = np.array([
    [1, 2, 3, 0],
    [0, 1, 2, 3],
    [3, 0, 1, 2],
    [2, 3, 0, 1]
])

filter_kernel = np.array([
    [2, 0, 1],
    [0, 1, 2],
    [1, 0, 2]
])

# Manual convolution calculation
def calculate_convolution(input_matrix, kernel, padding=0, stride=1):
    # Calculate output dimensions
    input_height, input_width = input_matrix.shape
    kernel_height, kernel_width = kernel.shape
    
    output_height = (input_height - kernel_height + 2 * padding) // stride + 1
    output_width = (input_width - kernel_width + 2 * padding) // stride + 1
    
    # Create output matrix
    output = np.zeros((output_height, output_width))
    
    # Perform convolution calculation
    for i in range(output_height):
        for j in range(output_width):
            # Determine current window position
            start_i = i * stride
            start_j = j * stride
            end_i = start_i + kernel_height
            end_j = start_j + kernel_width
            
            # Extract window and calculate dot product
            input_window = input_matrix[start_i:end_i, start_j:end_j]
            output[i, j] = np.sum(input_window * kernel)
    
    return output

# Calculate convolution result (no padding, stride=1)
convolution_result = calculate_convolution(input_data, filter_kernel)

# Get detailed computation process
def get_computation_details(input_matrix, kernel):
    details = []
    kernel_height, kernel_width = kernel.shape
    output_height, output_width = convolution_result.shape
    
    for i in range(output_height):
        for j in range(output_width):
            start_i, start_j = i, j
            end_i, end_j = i + kernel_height, j + kernel_width
            input_window = input_matrix[start_i:end_i, start_j:end_j]
            
            # Build calculation expression
            terms = []
            for ki in range(kernel_height):
                for kj in range(kernel_width):
                    input_val = input_window[ki, kj]
                    kernel_val = kernel[ki, kj]
                    if input_val != 0 and kernel_val != 0:  # Only show non-zero terms
                        terms.append("%d*%d" % (input_val, kernel_val))
                    
            computation = " + ".join(terms)
            result = convolution_result[i, j]
            details.append({
                "position": (i, j),
                "window": input_window,
                "computation": computation,
                "result": result
            })
    
    return details

# Get computation details
computation_details = get_computation_details(input_data, filter_kernel)

# Print convolution calculation process and result
print("CONVOLUTION EXAMPLE EXPLANATION")
print("=============================")
print("Input data:")
print(input_data)
print("\nFilter kernel:")
print(filter_kernel)
print("\nConvolution result:")
print(convolution_result)
print("\nDetailed calculation process:")

for detail in computation_details:
    i, j = detail["position"]
    print("\nOutput point (%d, %d):" % (i+1, j+1))
    print("Input window:")
    print(detail['window'])
    print("Calculation: %s = %d" % (detail['computation'], detail['result']))

# Create visualization
def visualize_convolution():
    fig, axs = plt.subplots(2, 2, figsize=(12, 10))
    
    # Input data
    axs[0, 0].matshow(input_data, cmap='gray', vmin=0, vmax=3)
    axs[0, 0].set_title('Input Data')
    for i in range(input_data.shape[0]):
        for j in range(input_data.shape[1]):
            axs[0, 0].text(j, i, str(input_data[i, j]), va='center', ha='center', fontsize=12)
    axs[0, 0].grid(True, color='black', linestyle='-', linewidth=0.5)
    
    # Filter kernel
    axs[0, 1].matshow(filter_kernel, cmap='gray', vmin=0, vmax=3)
    axs[0, 1].set_title('Filter Kernel')
    for i in range(filter_kernel.shape[0]):
        for j in range(filter_kernel.shape[1]):
            axs[0, 1].text(j, i, str(filter_kernel[i, j]), va='center', ha='center', fontsize=12)
    axs[0, 1].grid(True, color='black', linestyle='-', linewidth=0.5)
    
    # Convolution process example
    axs[1, 0].matshow(input_data, cmap='gray', vmin=0, vmax=3)
    axs[1, 0].set_title('Convolution Process Example')
    # Highlight first convolution window
    rect = Rectangle((-0.5, -0.5), 3, 3, fill=False, edgecolor='red', linewidth=2)
    axs[1, 0].add_patch(rect)
    for i in range(input_data.shape[0]):
        for j in range(input_data.shape[1]):
            axs[1, 0].text(j, i, str(input_data[i, j]), va='center', ha='center', fontsize=12)
    axs[1, 0].grid(True, color='black', linestyle='-', linewidth=0.5)
    
    # Convolution result
    axs[1, 1].matshow(convolution_result, cmap='viridis')
    axs[1, 1].set_title('Convolution Result')
    for i in range(convolution_result.shape[0]):
        for j in range(convolution_result.shape[1]):
            axs[1, 1].text(j, i, str(int(convolution_result[i, j])), va='center', ha='center', fontsize=12, color='white')
    axs[1, 1].grid(True, color='white', linestyle='-', linewidth=0.5)
    
    # Add convolution operator and arrow using ASCII
    fig.text(0.5, 0.5, '*', fontsize=30, ha='center', va='center')
    fig.text(0.75, 0.5, '->', fontsize=20, ha='center', va='center')
    
    # Adjust layout
    plt.tight_layout()
    plt.savefig('convolution_example_visualization.png', dpi=300, bbox_inches='tight')
    plt.close()

# Create detailed computation visualization
def visualize_computation_details():
    fig, axs = plt.subplots(2, 2, figsize=(15, 15))
    axs = axs.flatten()
    
    for idx, detail in enumerate(computation_details):
        i, j = detail["position"]
        ax = axs[idx]
        
        # Plot input data matrix
        ax.matshow(input_data, cmap='gray', vmin=0, vmax=3)
        
        # Highlight current computation window
        rect = Rectangle((j-0.5, i-0.5), 3, 3, fill=False, edgecolor='red', linewidth=2)
        ax.add_patch(rect)
        
        # Add value labels
        for row in range(input_data.shape[0]):
            for col in range(input_data.shape[1]):
                ax.text(col, row, str(input_data[row, col]), va='center', ha='center', fontsize=10)
        
        # Add computation labels at filter positions
        kernel_height, kernel_width = filter_kernel.shape
        for ki in range(kernel_height):
            for kj in range(kernel_width):
                input_val = input_data[i+ki, j+kj]
                kernel_val = filter_kernel[ki, kj]
                if input_val != 0 and kernel_val != 0:  # Only mark non-zero terms
                    ax.text(j+kj, i+ki, "%d*%d" % (input_val, kernel_val), 
                            va='bottom', ha='center', fontsize=8, color='red')
        
        # Add computation process and result below the figure
        ax.text(1.5, 4.5, "Calculation: %s = %d" % (detail['computation'], detail['result']), 
                va='center', ha='center', fontsize=10, 
                bbox=dict(boxstyle="round,pad=0.5", fc="yellow", ec="black", alpha=0.5))
        
        ax.set_title('Output Point (%d, %d) Calculation' % (i+1, j+1))
        ax.grid(True, color='black', linestyle='-', linewidth=0.3)
    
    plt.tight_layout()
    plt.savefig('convolution_computation_details.png', dpi=300, bbox_inches='tight')
    plt.close()

# Execute visualization
visualize_convolution()
visualize_computation_details()

# Display convolution formula explanation
print("\nCONVOLUTION FORMULA EXPLANATION:")
print("The convolution operation (f * g)(t) = Integral[f(tau) * g(t-tau)] dtau applied in discrete case:")
print("1. In image/matrix convolution, we use the 2D form: (I * K)(i,j) = Sum[I(m,n) * K(i-m,j-n)]")
print("2. For this example, we have a 4x4 input matrix I and a 3x3 convolution kernel K")
print("3. The convolution result R(i,j) is the sum of element-wise products between K and a 3x3 region of I")
print("4. Without padding and with stride=1, the output size is 2x2")
print("\nImages saved as: 'convolution_example_visualization.png' and 'convolution_computation_details.png'")